from langchain.prompts import PromptTemplate
from typing import List, Optional
from langchain.tools.base import BaseTool
from langchain.schema.output_parser import BaseOutputParser
from langchain.output_parsers import PydanticOutputParser
from langchain.prompts import load_prompt
import os, json
import tempfile

from langchain_core.prompts import PipelinePromptTemplate, BasePromptTemplate

# from AutoAgent.Action import Action

"""通过确保正确的字符编码使 JSON 字符串更易于中国用户阅读。"""
def chinese_friendly(string) -> str:
    lines = string.split('\n')
    """for i, line in enumerate(lines):：使用 enumerate 函数遍历 lines 列表，获取每一行的索引 i 和行内容 line。"""
    for i, line in enumerate(lines):
        if line.startswith('{') and line.endswith('}'):
            try:
                """尝试将当前行的内容加载为 JSON，并使用 json.dumps 将其重新转换为 JSON 字符串。
                json.loads(line) 将 JSON 字符串解析为 Python 对象，
                然后 json.dumps(..., ensure_ascii=False) 将 Python 对象转换为 JSON 字符串，
                其中 ensure_ascii=False 表示允许非 ASCII 字符出现在输出中。修改后的 JSON 字符串替换了原始的行内容。"""
                lines[i] = json.dumps(json.loads(line), ensure_ascii=False)
            except:
                pass
    return '\n'.join(lines)

"""将文件的内容加载为字符串。"""
def load_file(filename: str) -> str:
    """Loads a file into a string."""
    if not os.path.exists(filename):
        raise FileNotFoundError(f"File {filename} not found.")
    f = open(filename, 'r', encoding='utf-8')
    s = f.read()
    f.close()
    return s

"""构建主 prompt 模板的流程文件"""
class PromptTemplateBuilder:
    def __init__(
            self,
            prompt_path: str,
            prompt_file: str,
    ):
        self.prompt_path = prompt_path
        self.prompt_file = prompt_file

    """创建一个列出可用工具及其描述的提示字符串
    举例输出：
    Tool1: This is the first tool, args json schema: {"arg1": "value1", "arg2": "value2"}
    
    """
    def _get_tools_prompt(self, tools: List[BaseTool]) -> str:
        tools_prompt = ""
        for i, tool in enumerate(tools):
            prompt = f"{i + 1}. {tool.name}: {tool.description}, \
                        args json schema: {json.dumps(tool.args, ensure_ascii=False)}\n"
            tools_prompt += prompt
        return tools_prompt
    """
    {
    "template_path": "/path/to/templates/sub_template.json",
    "other_config": "some_value"
    }
    这个是临时文件的样子，最后用的是这个临时文件      
    """
    def _check_or_redirect(self, prompt_file: str) -> str:
        """
        这个部分使用 open 函数打开指定路径的文件。
        prompt_file 是文件的路径，
        "r" 表示以只读模式打开文件，而
        encoding="utf-8" 表示使用 UTF-8 编码读取文件内容。
        as f: 将打开的文件对象赋值给变量 f，这样我们可以通过 f 来操作文件。
        config = json.load(f): 使用 json.load 方法从文件对象 f 中加载 JSON 数据。
        这个方法会读取文件内容，并将其解析为 Python 字典或列表，然后将结果赋值给变量 config。
        """
        with open(prompt_file, "r", encoding='utf-8') as f:
            config = json.load(f)
        if "template_path" in config:
            # 如果是相对路径，则转换为绝对路径
            if not os.path.isabs(config["template_path"]):
                config["template_path"] = os.path.join(self.prompt_path, config["template_path"])
                # 生成临时文件
                tmp_file = tempfile.NamedTemporaryFile(
                    suffix='.json',
                    mode="w",
                    encoding="utf-8",
                    delete=False
                )
                tmp_file.write(json.dumps(config, ensure_ascii=False))
                tmp_file.close()
                return tmp_file.name
        return prompt_file

    def build(
            self,
            tools: Optional[List[BaseTool]] = None,
            output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:
        """将主提示模板文件的名称转换为绝对路径。
        os.path.join 是路径和文件名称的拼接，拼接为完整的文件路径"""
        main_file = os.path.join(self.prompt_path, self.prompt_file)
        """加载主提示模板文件。"""
        main_prompt_template = load_prompt(
            self._check_or_redirect(main_file)
        )
        """获取主提示模板中的输入变量列表。"""
        variables = main_prompt_template.input_variables
        """创建一个用于存储变量值的字典。"""
        partial_variables = {}
        """创建一个用于存储嵌套提示模板的列表。"""
        recursive_templates = []

        # 遍历所有变量，检查是否存在对应的模板文件
        for var in variables:
            # 是否存在嵌套模板
            if os.path.exists(os.path.join(self.prompt_path, f"{var}.json")):
                sub_template = PromptTemplateBuilder(
                    self.prompt_path, f"{var}.json"
                ).build(tools=tools, output_parser=output_parser)
                recursive_templates.append((var, sub_template))
            # 是否存在文本文件
            elif os.path.exists(os.path.join(self.prompt_path, f"{var}.txt")):
                var_str = load_file(os.path.join(self.prompt_path, f"{var}.txt"))
                partial_variables[var] = var_str

        if tools is not None and "tools" in variables:
            tools_prompt = self._get_tools_prompt(tools)
            partial_variables["tools"] = tools_prompt

        if output_parser is not None and "format_instructions" in variables:
            partial_variables["format_instructions"] = chinese_friendly(
                output_parser.get_format_instructions()
            )

        if recursive_templates:
            # 将有值嵌套的模板填充到主模板中
            main_prompt_template = PipelinePromptTemplate(
                final_prompt=main_prompt_template,
                pipeline_prompts=recursive_templates
            )

        # 将有值的变量填充到模板中
        main_prompt_template = main_prompt_template.partial(**partial_variables)

        return main_prompt_template


if __name__ == "__main__":
    builder = PromptTemplateBuilder("../prompts/main", "main.json")
    # output_parser = PydanticOutputParser(pydantic_object=Action)
    prompt_template = builder.build(tools=[], output_parser=None)
    print(prompt_template.format(
        ai_name="瓜瓜",
        ai_role="智能助手机器人",
        task_description="解决问题",
        work_dir=".",
        short_term_memory="",
        long_term_memory="",
        tools="",
         format_instructions=""
    ))
